import torch
import torchvision.transforms as transforms
from PIL import Image
import torch.nn.functional as F
from network import AlexNet


def main():
    with open('../../datasets/flowers/classes.txt', 'r') as text_cls:
        classes_str = text_cls.read()
    classes = classes_str.split(',')   # classes = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']

    transform = transforms.Compose([transforms.Resize((224, 224)),
                                    transforms.ToTensor(),
                                    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
    
    net = AlexNet(num_classes=5)
    net.load_state_dict(torch.load('../../models/AlexNet_best.pth', map_location='cpu')['state_dict'])

    im = Image.open('../../datasets/flowers/test/sunflowers/27465811_9477c9d044.jpg') 
    # print(np.array(im).shape)
    im = transform(im)  # [C, H, W]
    # print(np.array(im).shape)
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]
    # print(np.array(im).shape)
    
    net.eval()
    with torch.no_grad():
        outputs = net(im)
        print(outputs.shape, F.softmax(outputs, dim=1))
        predict = torch.max(outputs, dim=1)[1].numpy()
        print(classes[int(predict)])


if __name__ == '__main__':
    main()
